#!/usr/bin/env python
# encoding: utf-8
import numpy as np
import tensorflow as tf
import torch
from tensorflow.python.framework import ops
from tensorflow.python.keras import backend as K
from tensorflow.python.ops import math_ops

# tensorflow版本==========================================================================

def guassian_kernel(source, target, kernel_mul=2.0, kernel_num=5, fix_sigma=None):
    n_s = tf.shape(source)[0]
    n_s = 64 if n_s is None else n_s
    n_t = tf.shape(target)[0]
    n_t = 64 if n_t is None else n_t
    n_samples = n_s + n_t
    total = tf.concat([source, target], axis=0)  # [None,n]
    total0 = tf.expand_dims(total, axis=0)  # [1,b,n]
    total1 = tf.expand_dims(total, axis=1)  # [b,1,n]
    L2_distance = tf.reduce_sum(((total0 - total1) ** 2),
                                axis=2)  # [b,b,n]=>[b,b]                                 #   [None,None,n]=>[128,128,1]
    if fix_sigma:
        bandwidth = fix_sigma
    else:
        bandwidth = tf.reduce_sum(L2_distance) / float(n_samples ** 2 - n_samples)
    bandwidth /= kernel_mul ** (kernel_num // 2)
    bandwidth_list = [bandwidth * (kernel_mul ** i) for i in range(kernel_num)]
    kernel_val = [tf.exp(-L2_distance / bandwidth_temp) for bandwidth_temp in bandwidth_list]
    return sum(kernel_val)  # /len(kernel_val)   #[b,b]


def MMD(source, target, kernel_mul=2.0, kernel_num=5, fix_sigma=None):
    kernels = guassian_kernel(source, target,
                              kernel_mul=kernel_mul, kernel_num=kernel_num, fix_sigma=fix_sigma)
    n_s = tf.shape(source)[0]
    n_s = 64 if n_s is None else n_s
    n_t = tf.shape(target)[0]
    n_t = 64 if n_t is None else n_t
    XX = tf.reduce_sum(kernels[:n_s, :n_s]) / float(n_s ** 2)
    YY = tf.reduce_sum(kernels[-n_t:, -n_t:]) / float(n_t ** 2)
    XY = tf.reduce_sum(kernels[:n_s, -n_t:]) / float(n_s * n_t)
    YX = tf.reduce_sum(kernels[-n_t:, :n_s]) / float(n_s * n_t)
    loss = XX + YY - XY - YX
    return loss

def mmd_loss(y_true, y_pred):

    mmd = MMD(y_pred, y_true)
    y_pred = ops.convert_to_tensor_v2(y_pred)
    y_true = math_ops.cast(y_true, y_pred.dtype)
    return K.mean(math_ops.squared_difference(y_pred, y_true) + mmd, axis=-1)

if __name__ == '__main__':
    source = [[1], [2], [3], [4], [5]]
    source = np.array(source)
    source = tf.convert_to_tensor(source, dtype=tf.float32)
    target = [[3], [4], [5], [6], [7]]
    target = np.array(target)
    target = tf.convert_to_tensor(target, dtype=tf.float32)
    loss = MMD(source, target)
    print(loss)
